# coding=UTF-8
import numpy as np
import cv2
import os
import uuid
import logging
import pydicom
from abc import ABCMeta, abstractmethod
import SimpleITK as sitk
from scipy import stats

GLOBAL_CODEC = ['utf-8', 'iso-8859-1']
__TRACE__ = True if os.getenv('TRACE', 'True') == 'True' else False
UDOG_FLAG = False
VERSION = "1.6.6.1"
SOURCE_PATH = "source_path"
DATA_FORMAT = "data_format"
IMAGE_UID = "image_uid"
INSTANCE_NUMBERS = "instance_numbers"
SOP_INSTANCE_UIDS = "sop_instance_uids"
INSTANCE_POSITIONS = "instance_positions"
COLOR_GRAY = "GRAY"
EUCLIDEAN_PHYSICAL = "euclidean_physical_space"
EUCLIDEAN_PIXEL_SPACE = "euclidean_pixel_space"
GTD_EXTENSION = "gtd"
XML_EXTENSION = "xml"
TXT_EXTENSION = 'txt'
ROI_STR = "roi"
ROI_STR_POSITIVE = "+roi"
ROI_STR_NEGATIVE = "-roi"
OUTER_STR = "outer"
INNER_STR = "inner"
GLOBAL_CODEC = ['utf-8', 'iso-8859-1']
HBASE_COLUMN_FAMILY = "CF_DEFAULT"


class Serializable(object):
    __metaclass__ = ABCMeta
    
    @abstractmethod
    def to_bytes(self):
        raise NotImplementedError
    
    @abstractmethod
    def from_bytes(self, raw_bytes):
        raise NotImplementedError


class Image(Serializable):
    """ image object consists of two major components, a header describes common image
    attributes and a numpy nd-array storing image pixels."""
    magic_number = 0x11
    
    def __int__(self):
        self.pixel_data = None
        self.header = None
    
    # @abstractmethod
    def to_bytes(self):
        raise NotImplementedError
    
    # @abstractmethod
    def from_bytes(self, raw_bytes):
        raise NotImplementedError


class ImageHeader(object):
    """ the header object describe basic attributes regarding image pixels
    and how it is produced. """
    
    def __init__(self,
                 dim=None,
                 dtype=None,
                 stype=None,
                 channels=None,
                 color_space=None,
                 origin=None,
                 direction=None,
                 spacing=None,
                 size=None,
                 shape=None,
                 metadata=None,
                 slice_thickness=None
                 ):
        self.uid = str(uuid.uuid4().hex)
        self.dim = dim
        self.dtype = dtype
        self.stype = stype
        self.channels = channels
        self.color_space = color_space
        self.origin = origin
        self.direction = direction
        self.spacing = spacing
        self.size = size
        self.shape = shape
        self.metadata = metadata
        self.slice_thickness = slice_thickness
    
    def __cmp__(self, other):
        if (self.dim, self.dtype, self.stype, self.channels, self.color_space, self.origin,
            self.direction, self.spacing, self.size, self.shape, self.metadata
                == other.dim, other.dtype, other.stype, other.channels,
            self.color_space, self.origin
            , other.direction, other.spacing, other.size, other.shape, other.metadata):
            return 0
        return -1


def itk_zyx_to_yxz(imagearray):
    """SimpleItk stores pixels in order of width, height and depth, when it's generated from GetArrayFromImage
     sitk implies the order of pixels to be plane, rows, cols, so here we just flip plane index up-and-down,
     and swap z-plane with the rows, output array would be rows, cols, flipped(z-plane)"""
    if len(imagearray.shape) == 2:
        return imagearray
    imagearray = np.moveaxis(imagearray, 0, -1)
    return imagearray


def convert_from_itk(itkdata):
    itkoutput = sitk.GetArrayFromImage(itkdata)
    return itk_zyx_to_yxz(itkoutput)


def normalize(pixels, from_span=None, to_span=(0, 255)):
    p = pixels.astype(np.float32)
    if from_span is not None:
        p[p > from_span[1]] = from_span[1]
        p[p < from_span[0]] = from_span[0]
    else:
        from_span = (np.min(p), np.max(p))
    if from_span[0] == from_span[1]:
        from_span = (from_span[0], from_span[0] + 1)
    p = ((p - from_span[0]) / float(from_span[1] - from_span[0])) * (to_span[1] - to_span[0]) + to_span[0]
    return p.astype(np.uint8)


def convert_any_string_to_utf8(test_string, decoder_list):
    string_unicode = test_string.decode('utf-8', 'ignore')
    return string_unicode.encode('utf-8')


def fetch_metadata_from_itk_image(image):
    metadata_keys = image.GetMetaDataKeys()
    metadata = {}
    for k in metadata_keys:
        v = image.GetMetaData(k).strip()
        if len(v) != 0:
            key = k.strip()
            # metadata[convert_any_string_to_utf8(key, GLOBAL_CODEC)] \
            #     = convert_any_string_to_utf8(v, GLOBAL_CODEC)
            metadata[key] = v
    return metadata


def monochrome_convert(image, default_window=True):
    try:
        pminterp = image.header.metadata.get('0028|0004')
        
        if pminterp:
            pminterp = pminterp.strip().upper()
        
        pixels = image.pixel_data
        
        if pminterp == 'MONOCHROME1':
            bits_stored = image.header.metadata.get('0028|0101')
            if bits_stored:
                maxval = 2 ** int(bits_stored) - 1
            else:
                maxval = 4095  # A sane default (12-bit scanner)
            
            if default_window:
                try:
                    window_center = image.header.metadata.get("0028|1050").split('\\')[0]
                    window_width = image.header.metadata.get("0028|1051").split('\\')[0]
                    if window_center and window_width:
                        pixels[pixels > (int(window_center) + int(window_width) / 2)] = int(window_center) + int(
                            window_width) / 2
                        pixels[pixels < (int(window_center) - int(window_width) / 2)] = int(window_center) - int(
                            window_width) / 2
                except:
                    pass
            pixels = maxval - pixels
        
        elif pminterp == 'MONOCHROME2':
            if default_window:
                try:
                    window_center = image.header.metadata.get("0028|1050").split('\\')[0]
                    window_width = image.header.metadata.get("0028|1051").split('\\')[0]
                    
                    if window_center and window_width:
                        pixels[pixels > (int(window_center) + int(window_width) / 2)] = int(window_center) + int(
                            window_width) / 2
                        pixels[pixels < (int(window_center) - int(window_width) / 2)] = int(window_center) - int(
                            window_width) / 2
                except:
                    pass
        
        elif pminterp == 'RGB':
            colorweights = [0.299, 0.587, 0.114]
            pixels = np.mean(pixels * colorweights, axis=2)
    except:
        pixels = image.pixel_data
    
    if window_center and window_width:
        if isinstance(window_center, list):
            wl = int(window_center)
        else:
            wl = int(window_center)
        if isinstance(window_width, list):
            ww = int(window_width)
        else:
            ww = int(window_width)
        image = normalize(pixels, (wl - ww / 2, wl + ww / 2))
    else:
        image = normalize(pixels)
    image = np.concatenate((image, image, image), -1)
    # print(image.shape)
    return image


def read_dcm(old_image):
    image = monochrome_convert(old_image)
    
    window_center = old_image.metadata.get("WindowCenter")
    window_width = old_image.metadata.get("WindowWidth")
    if window_center and window_width:
        if isinstance(window_center, list):
            wl = window_center[0]
        else:
            wl = window_center
        if isinstance(window_width, list):
            ww = window_width[0]
        else:
            ww = window_width
        image = normalize(image, (wl - ww / 2, wl + ww / 2))
    else:
        image = normalize(image)
    image = (image.astype(np.uint8))[0]
    image = image.astype(np.uint8)
    image = np.stack((image, image, image), -1)
    return image


# def read_data(dir):
#     series_reader = sitk.ImageSeriesReader()
#     series_reader.SetOutputPixelType(sitk.sitkInt16)
#     series_reader.SetFileNames(series_file_names)
#     series_image = series_reader.Execute()
#     image = Image()
#     image_header = ImageHeader()


def read_dcm_from_dir(image_dir, dicom_sorted=None):
    logger = logging.getLogger(__name__)
    instance = pydicom.dcmread(image_dir, stop_before_pixels=True, force=True)
    ds = sitk.ReadImage(image_dir)
    metadata_dict = fetch_metadata_from_itk_image(ds)
    series_instance_uid = instance.SeriesInstanceUID.strip()
    sop_instance_uids = instance.SOPInstanceUID.strip()
    try:
        instance_positions = [float(instance.ImagePositionPatient[-1])]
    except Exception:
        instance_positions = []
    series_reader = sitk.ImageFileReader()
    series_reader.SetOutputPixelType(sitk.sitkInt16)
    series_reader.SetFileName(image_dir)
    series_image = series_reader.Execute()
    image = Image()
    image_header = ImageHeader()
    image_header.uid = series_instance_uid
    image_header.metadata = metadata_dict
    image_header.metadata[SOURCE_PATH] = os.path.dirname(image_dir)
    image_header.metadata[DATA_FORMAT] = 'dcm'
    image_header.metadata[INSTANCE_NUMBERS] = 1
    image_header.metadata[SOP_INSTANCE_UIDS] = sop_instance_uids
    image_header.metadata[INSTANCE_POSITIONS] = instance_positions
    image_header.channels = 1
    image_header.dim = series_image.GetDimension()
    image_header.color_space = COLOR_GRAY
    image_header.dtype = '<i2'
    image_header.stype = EUCLIDEAN_PHYSICAL
    image_header.size = list(series_image.GetSize())
    image_header.size[0], image_header.size[1] = image_header.size[1], image_header.size[0]
    image_header.origin = list(series_image.GetOrigin())
    image_header.origin[0], image_header.origin[1] = \
        image_header.origin[1], image_header.origin[0]
    image_header.direction = series_image.GetDirection()
    image_header.direction = np.array(image_header.direction) \
        .reshape(image_header.dim, image_header.dim).tolist()
    image_header.direction[0], image_header.direction[1] = \
        image_header.direction[1], image_header.direction[0]
    image_header.spacing = list(series_image.GetSpacing())
    image_header.spacing[0], image_header.spacing[1] = \
        image_header.spacing[1], image_header.spacing[0]
    image_header.slice_thickness = None
    image.pixel_data = convert_from_itk(series_image).astype(image_header.dtype)
    image_header.shape = image.pixel_data.shape
    image.header = image_header
    return image, image_header.spacing[0]


